﻿#include "utils.h"
#include <algorithm>

int main(int argc, char* argv[]) {


 	// ULONG_PTR PEBaseAddress = getPEBaseAddr(L"D:/c++/code/p_shell/Release/load_poc.dll");
	/*
	ULONG_PTR PEBaseAddress = getPEBaseAddr(L"D:/c++/code/p_shell/Release/load_poc_exe.exe");
	
	if (PEBaseAddress)
	{
		ReflectiveLoader(PEBaseAddress);
	}
	*/

	//printf("%p", PEBaseAddress);

/*
	std::string filePath = "D:/c++/code/p_shell/Release/load_poc_exe.exe"; // 要加密的文件路径
	std::string key = "hello";        // 异或加密使用的密钥

	xorEncryptDecrypt(filePath, key);

	std::cout << TEXT("文件已加密并保存为 ") << filePath << ".xor" << std::endl;

	std::cout << "Test encoding" << std::endl;
	printf("Test encoding");
*/
	PCWSTR pszLibFile = NULL;
	char* strProcName;
	DWORD dwProcessId = 0;
	DWORD dwTechnique = 0;

	if (argc != 3)
	{
		displayHelp();
		return(0);
	}

	std::string key = "hello";        // 异或加密使用的密钥


	// 寻找最后一个反斜杠（'\'）的位置
	// std::string fullPath(argv[2]);
	// size_t lastBackslashPos = fullPath.find_last_of("\\");

	// 提取目录路径
	// std::string directoryPath = fullPath.substr(0, lastBackslashPos + 1);

	if (strcmp(argv[1], "-e") == 0)
	{
		strProcName = (char*)malloc((strlen(argv[2]) + 1) * sizeof(char));
		strcpy_s(strProcName, (strlen(argv[2]) + 1) * sizeof(char), argv[2]);
		xorEncryptDecrypt((const std::string&)strProcName,key, 0);
		printf("Shell addition completed: %s%s", argv[2], ".shell");
	}
	else if (strcmp(argv[1], "-d") == 0 )
	{
		strProcName = (char*)malloc((strlen(argv[2]) + 1) * sizeof(char));
		strcpy_s(strProcName, (strlen(argv[2]) + 1) * sizeof(char), argv[2]);
		xorEncryptDecrypt((const std::string&)strProcName, key, 1);
		printf("Shell removal completed: %s%s", argv[2], ".exe");
      
		std::string filePath(argv[2]);
		std::string  newFilePath = filePath + ".exe";

		// 调用MultiByteToWideChar以确定转换后的宽字符字符串长度
		int wcharCount = MultiByteToWideChar(CP_UTF8, 0, newFilePath.c_str(), -1, nullptr, 0);

		// 为宽字符字符串分配内存（包含null终止符）
		wchar_t* wcharArray = new wchar_t[wcharCount];

		// 执行从UTF-8到UTF-16的转换
		MultiByteToWideChar(CP_UTF8, 0, newFilePath.c_str(), -1, wcharArray, wcharCount);


		// 将 PE 文件在内存中进行分析、展开、获取入口函数执行

		ULONG_PTR PEBaseAddress = getPEBaseAddr(wcharArray);

		if (PEBaseAddress)
		{
			ReflectiveLoader(PEBaseAddress);
		}


	}
	else
	{
		displayHelp();
	}

	system("pause");
	return 0;
}